package BinaryTree.Medium;

import Structures.BinaryTree;
import Structures.TreeNode;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;
import java.util.Stack;

public class LC0437 {
    /**
     * 我的解法，思路是保存根结点到结点p(k)的路径上，从各个结点开始到p(k)的和。
     * 时间复杂度和空间复杂度都高于官方解法。
     */
    public int pathSum(TreeNode root, int targetSum) {
        if (root == null) return 0;

        Stack<TreeNode> nodeStack = new Stack<>();
        Stack<ArrayList<Long>> sumTraceStack = new Stack<>();
        nodeStack.push(root);
        ArrayList<Long> motherTrace = new ArrayList<>();
        sumTraceStack.push(motherTrace);

        int nPaths = 0;
        while (!nodeStack.empty()) {
            TreeNode node = nodeStack.pop();
            ArrayList<Long> trace = sumTraceStack.pop();
            for (int i = 0; i < trace.size(); i++) {
                long pathSum = trace.get(i) + node.val;
                if (pathSum == targetSum) nPaths++;
                trace.set(i, pathSum);
            }
            if (node.val == targetSum) nPaths++;
            trace.add((long) node.val);
            if (node.right != null) {
                nodeStack.add(node.right);
                sumTraceStack.add((ArrayList<Long>) trace.clone());
            }
            if (node.left != null) {
                nodeStack.add(node.left);
                sumTraceStack.add((ArrayList<Long>) trace.clone());
            }
        }

        return nPaths;
    }

    /**
     * 官方题解，基于前缀和。具体思路请参考官方题解。
     * 某个节点i的“前缀和”定义为从根结点到i的路径上，包含根结点、不包含i的所有结点的值之和。
     */
    public static int pathSumPrefix(TreeNode root, int targetSum) {
        Map<Long, Integer> prefix = new HashMap<>(); // 前缀和->路径数量映射
        prefix.put(0L, 1); // 如果题目所要求的“路径”包含根结点，则需要一个key=0的项来确保计算正确
        return dfs(root, prefix, 0, targetSum);
    }

    /**
     * 返回包含该子树中结点的“目标路径”的总数
     *
     * @param curr 根结点到当前结点之前的一个结点的和
     */
    public static int dfs(TreeNode node, Map<Long, Integer> prefix, long curr, int targetSum) {
        if (node == null) return 0;

        curr += node.val; // 根结点到该结点位置的和（包含该结点）
        // 假设当前结点为p(k)，如果根结点到p(k)的和为curr，且根结点到p(i-1)的和为curr-target，那么p(i)到p(k)的和就为target。
        int ret = prefix.getOrDefault(curr - targetSum, 0);
        prefix.put(curr, prefix.getOrDefault(curr, 0) + 1);
        ret += dfs(node.left, prefix, curr, targetSum);
        ret += dfs(node.right, prefix, curr, targetSum);
        prefix.put(curr, prefix.get(curr) - 1); // 恢复状态，当前结点的子结点处理完后，应删除依赖于当前结点的前缀和信息
        return ret;
    }

    public static void main(String[] args) {
        Integer[] array = new Integer[]{10, 5, -3, 3, 2, null, 11, 3, -2, null, 1};
        TreeNode root = BinaryTree.constructBinaryTree(array);
        System.out.println(pathSumPrefix(root, 8));
    }
}
